import os
import re
from collections import defaultdict
from PIL import Image, ImageDraw
import shutil

def blend_images(base_img, eff_imgs):
    """
    将基础图片和效果图片合成
    
    Args:
        base_img (Image): 基础图片
        eff_imgs (list): 效果图片列表
    
    Returns:
        Image: 合成后的图片
    """
    # 确保所有图片大小一致
    base_size = base_img.size
    
    # 调整所有效果图片大小以匹配基础图片
    resized_eff_imgs = []
    for eff_img in eff_imgs:
        if eff_img.size != base_size:
            eff_img = eff_img.resize(base_size, Image.LANCZOS)
        resized_eff_imgs.append(eff_img)
    
    # 将基础图片转换为RGBA模式
    base_img = base_img.convert("RGBA")
    
    # 对每个效果图片进行合成
    result_img = base_img.copy()
    for eff_img in resized_eff_imgs:
        eff_img = eff_img.convert("RGBA")
        # 使用alpha混合将效果图片叠加到结果图片上
        result_img = Image.alpha_composite(result_img, eff_img)
    
    return result_img

def extract_id_from_filename(filename):
    """
    从文件名中提取ID用于匹配
    例如: nt002_AT01_M01_001.png -> 001
    """
    # 匹配文件名末尾的数字（在最后一个下划线后）
    match = re.search(r'_([0-9]+)\.png$', filename)
    if match:
        return match.group(1)
    return None

def find_eff_directories(images_dir):
    """
    查找images目录下所有以eff开头的目录
    
    Args:
        images_dir (str): images目录路径
    
    Returns:
        list: eff目录路径列表
    """
    eff_dirs = []
    for item in os.listdir(images_dir):
        item_path = os.path.join(images_dir, item)
        if os.path.isdir(item_path) and item.startswith('eff'):
            eff_dirs.append(item_path)
    return eff_dirs

def copy_or_compose_images(base_dir, images_dir, output_dir):
    """
    将base目录下的所有图片处理后输出到output目录
    自动查找所有eff目录并使用其中的效果图片
    有匹配效果图片的进行合成，没有的直接复制
    统一使用composed_前缀命名所有输出图片
    
    Args:
        base_dir (str): 基础图片目录
        images_dir (str): images目录路径
        output_dir (str): 输出图片目录
    """
    # 查找所有eff目录
    eff_dirs = find_eff_directories(images_dir)
    print(f"找到 {len(eff_dirs)} 个效果图片目录: {eff_dirs}")
    
    # 获取所有效果图片
    eff_files = defaultdict(list)
    for eff_dir in eff_dirs:
        print(f"处理目录: {eff_dir}")
        for file in os.listdir(eff_dir):
            if file.endswith('.png'):
                file_id = extract_id_from_filename(file)
                if file_id:
                    eff_files[file_id].append(os.path.join(eff_dir, file))
                    print(f"  发现效果图片: {file} (ID: {file_id})")
    
    # 处理所有基础图片
    count = 0
    composed_count = 0
    copied_count = 0
    
    # 先列出所有基础图片
    base_files = []
    for file in os.listdir(base_dir):
        if file.endswith('.png'):
            base_files.append(file)
    
    # 按文件名排序以确保处理顺序
    base_files.sort()
    
    for file in base_files:
        file_id = extract_id_from_filename(file)
        base_path = os.path.join(base_dir, file)
        
        print(f"处理基础图片: {file} (ID: {file_id})")
        
        if file_id and file_id in eff_files:
            # 有匹配的效果图片，进行合成
            try:
                # 加载基础图片
                base_img = Image.open(base_path)
                
                # 加载效果图片
                eff_imgs = []
                for eff_path in eff_files[file_id]:
                    eff_img = Image.open(eff_path)
                    eff_imgs.append(eff_img)
                
                # 合成图片
                composed_img = blend_images(base_img, eff_imgs)
                
                # 生成输出文件名（统一使用composed_前缀）
                output_filename = f"composed_{file_id}.png"
                output_path = os.path.join(output_dir, output_filename)
                
                # 保存合成结果
                composed_img.save(output_path, 'PNG')
                
                print(f"  已合成图片: {output_path} (使用 {len(eff_imgs)} 个效果图片)")
                composed_count += 1
            except Exception as e:
                print(f"  合成 {file} 时出错: {e}")
        else:
            # 没有匹配的效果图片，复制并重命名
            try:
                # 生成输出文件名（统一使用composed_前缀）
                output_filename = f"composed_{file_id}.png"
                output_path = os.path.join(output_dir, output_filename)
                
                # 复制文件
                shutil.copy2(base_path, output_path)
                
                print(f"  已复制图片: {output_path}")
                copied_count += 1
            except Exception as e:
                print(f"  复制 {file} 时出错: {e}")
        
        count += 1
    
    print(f"\n总共处理 {count} 张图片:")
    print(f"  - 合成图片: {composed_count} 张")
    print(f"  - 复制图片: {copied_count} 张")

def main():
    # 定义目录
    base_dir = 'images/base'
    images_dir = 'images'
    output_dir = 'output'
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 处理图片
    copy_or_compose_images(base_dir, images_dir, output_dir)

if __name__ == "__main__":
    main()